import os
import shutil


def separate_files(source_dir, png_target_dir, txt_target_dir):
    """
    将 CAISA-B 数据集中的 PNG 文件和 TXT 文件分开到两个目录中，保持原始文件夹结构。

    Args:
        source_dir (str): 源数据集的根目录路径。
        png_target_dir (str): 用于存放 PNG 文件的目标目录路径。
        txt_target_dir (str): 用于存放 TXT 文件的目标目录路径。
    """
    # 遍历源目录
    for root, dirs, files in os.walk(source_dir):
        # 计算当前路径相对于源目录的相对路径
        relative_path = os.path.relpath(root, source_dir)

        # 创建目标目录（PNG 和 TXT）
        png_dir = os.path.join(png_target_dir, relative_path)
        txt_dir = os.path.join(txt_target_dir, relative_path)

        # 确保目标目录存在
        os.makedirs(png_dir, exist_ok=True)
        os.makedirs(txt_dir, exist_ok=True)

        # 遍历当前目录下的文件
        for file in files:
            file_path = os.path.join(root, file)

            # 判断文件类型
            if file.lower().endswith('.png'):
                # 复制 PNG 文件到 PNG 目标目录
                shutil.copy2(file_path, png_dir)
                print(f"Copied PNG: {file_path} -> {png_dir}")
            elif file.lower().endswith('.txt'):
                # 复制 TXT 文件到 TXT 目标目录
                shutil.copy2(file_path, txt_dir)
                print(f"Copied TXT: {file_path} -> {txt_dir}")


if __name__ == "__main__":
    # 设置源目录和目标目录
    source_directory = "E:/datasets/CAISA-B"  # 替换为你的 CAISA-B 数据集路径
    png_target_directory = "E:/datasets/CAISA-B-png"  # PNG 文件的目标目录
    txt_target_directory = "E:/datasets/CAISA-B-txt"  # TXT 文件的目标目录

    # 执行文件分离
    separate_files(source_directory, png_target_directory, txt_target_directory)
    print("File separation completed!")